package ace.cmp.spring.graphql.data.jdbc.querydsl.core.dataloader.impl;

import ace.cmp.spring.graphql.core.execution.AceBatchLoaderRegistry;
import ace.cmp.spring.graphql.data.jdbc.querydsl.core.dataloader.GraphqlControllerDataLoaderRegister;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import reactor.core.publisher.Mono;

/**
 * @author caspar
 * @date 2023/9/26 15:02
 */
public class GraphqlControllerDataLoaderRegisterImpl
    implements GraphqlControllerDataLoaderRegister {
  private final AceBatchLoaderRegistry loaderRegistry;

  public GraphqlControllerDataLoaderRegisterImpl(AceBatchLoaderRegistry loaderRegistry) {
    this.loaderRegistry = loaderRegistry;
  }

  @Override
  public <TId, TIEntityDto, TEntityDto extends TIEntityDto> void register(
      Class<TId> idClass,
      Class<TIEntityDto> iEntityClass,
      Function<List<TId>, Map<TId, TEntityDto>> getMapByIdsInFn) {
    if (loaderRegistry.hasMapperLoader(iEntityClass.getName())) {
      return;
    }
    loaderRegistry
        .forTypePair(idClass, iEntityClass)
        .registerMappedBatchLoader(
            (ids, env) -> {
              return Mono.fromSupplier(
                  () -> {
                    Map<TId, TEntityDto> entities = getMapByIdsInFn.apply(ids.stream().toList());
                    Map<TId, TIEntityDto> dtoMap =
                        entities.entrySet().stream()
                            .collect(
                                HashMap::new,
                                (m, entry) -> {
                                  TId id = entry.getKey();
                                  TIEntityDto entityDto = entry.getValue();
                                  m.put(id, entityDto);
                                },
                                HashMap::putAll);
                    return dtoMap;
                  });
            });
  }
}
